#coding:utf-8
import numpy as np
import cv2
from transforms3d import quaternions
import math
from scipy.optimize import least_squares

def method_old():
    pose1 = "332.800000 -0.028192952 -0.110917628 0.307533890 -0.001780688 -0.170053825 -0.567205548 0.805826545"
    pose2 = "363.000000 -0.035459653 -0.121275246 0.323589832 -0.074926689 -0.076565869 0.119611472 0.987024188"

    c1 = np.array([421,180,1])
    w1 = math.sqrt((409-433)**2 + (169-169)**2)
    c2 = np.array([357,249,1])
    w2 = math.sqrt((401-392)**2 + (217-291)**2)
    t_w = 60

    pose1 = pose1.split(' ')[1:]
    pose2 = pose2.split(' ')[1:]
    t1 = np.array([float(x) for x in pose1[:3]])
    q1 = np.array([float(x) for x in pose1[3:]])
    R1 = quaternions.quat2mat(q1)
    t2 = np.array([float(x) for x in pose2[:3]])
    q2 = np.array([float(x) for x in pose2[3:]])
    R2 = quaternions.quat2mat(q2)

    R21 = R2.transpose().dot(R1)
    T21 = R2.transpose().dot(t1) - R2.transpose().dot(t2)

    K = np.array([[391,0,320],[0,391,180],[0,0,1]])
    XT1 = np.linalg.inv(K).dot(c1)*(K[0,0]/w1*t_w)
    XT2 = np.linalg.inv(K).dot(c2)*(K[0,0]/w2*t_w)

    tmp = XT2 - R21.dot(XT1)
    scale = np.sum(T21*tmp) / np.sum(T21*T21)

    scale = 7000
    Off = R1.dot(XT1) + t1*scale

    print Off
    print scale

def loss_fun(x, Rcw, Tcw, pix, K):
    Twq = x[:3]
    scale = x[-1]
    uhead = np.zeros(len(Rcw)*2,)
    for k in range(len(Rcw)):
        Xc = Rcw[k].dot(Twq) + scale*Tcw[k]
        x = Xc / Xc[-1]
        uhead[k*2:(k+1)*2] = K.dot(x)[:2]
    return uhead - pix

def method_opt():
    data_lines = open('get_slam_params_data2.txt').readlines()
    # pose in data_lines are current to world, Rwc
    Rcw = []
    Tcw = []
    pix = []
    for line in data_lines:
        data = line.split(' ')[1:]
        # data 1
        #Twc = np.array([float(x) for x in data[:3]])
        #q = np.array([float(x) for x in data[3:7]])
        #Rwc = quaternions.quat2mat(q)
        #Rcw.append(Rwc.transpose())
        #Tcw.append(-Rwc.transpose().dot(Twc))
        # data 2
        Tcw.append(np.array([float(x) for x in data[:3]]))
        Rcw.append(cv2.Rodrigues(np.array([float(x) for x in data[3:6]]))[0])
        #pix.append(np.array([float(x) for x in data[7:]]))
        pix.append(np.array([float(x) for x in data[6:]]))
    pix = np.array(pix).flatten()
    x0 = np.array([-500,-400,4000,12000])
    K = np.array([[391,0,320],[0,391,180],[0,0,1]])
    res = least_squares(loss_fun, x0,args=(Rcw, Tcw, pix, K))
    print 'res = ', res.x
    print 'cost = ', res.cost
    
if __name__ == "__main__":
    method_opt()